
#include "triu_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

  TriuTilingData tiling;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    auto coreNum = ascendcPlatform.GetCoreNum();
    auto BLOCK_DIM = 1;
    context->SetBlockDim(BLOCK_DIM);
    
    const auto inputDataType=context->GetInputTensor(0)->GetDataType();
    uint32_t typeSize=0;
    uint64_t key=1;
    //buffer for queue
	uint64_t UB_SHARING_NUM=2;
    switch(inputDataType){
        case ge::DT_FLOAT:
            typeSize=sizeof(float);
            break;
        case ge::DT_FLOAT16:
            typeSize=2;
	        break;
        default:
            //ASSERT(0 && "Unknown inputDataType");
            typeSize=4;
            break;
    }

    const auto inputShape=context->GetInputTensor(0)->GetOriginShape();
    //class Shape: size_t dim_num_; int64_t dims_[];
    int64_t dimSize=inputShape.GetDimNum();
    int64_t matrixNum=1;
    int64_t matrixSize=1;
    int i;
    for(i=0;i<dimSize-2;i++){
        matrixNum*=inputShape.GetDim(i);
    }
    int64_t rowLength=inputShape.GetDim(i);
    i++;
    int64_t columnLength=inputShape.GetDim(i);
    matrixSize=rowLength*columnLength;

    const auto runtime_attrs=context->GetAttrs();
    const int64_t* diagPtr=runtime_attrs->GetInt(0);
    int64_t diagVal=*diagPtr;
    if(diagVal<columnLength-1 && diagVal>-rowLength){
        //regular
        key=1;
    }else if(diagVal<=-rowLength){
        //itself, TQueBind is enough
        key=2;
    }else{
        //all zero, just copyIn, Sub and copyOut
        key=3;
    }
    uint32_t ALIGN_NUM = 32 / typeSize;
    uint32_t totalLengthAligned = (matrixNum * matrixSize + ALIGN_NUM - 1) / ALIGN_NUM * ALIGN_NUM;
    //in an individual core
    uint64_t ub_size=0;
    uint64_t loopCnt=0;
    uint64_t fullTileLength = 0;
    uint64_t lastTileLength = 0;
    int32_t fullCnt=0,lastCnt=0;
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
    uint64_t ub_length = ( (ub_size / typeSize / UB_SHARING_NUM) / ALIGN_NUM * ALIGN_NUM ) - ALIGN_NUM;
    if(key==1 && diagVal<=0 && columnLength % (256/typeSize)==0){
        //只处理对齐的，而且每次只处理一个大矩阵里的几行
        key=4;
        UB_SHARING_NUM=4;
        ub_length = ( (ub_size)  / typeSize / UB_SHARING_NUM) / columnLength * columnLength;
        loopCnt = (matrixSize + ub_length - 1) / ub_length;
        if(loopCnt==1){
            fullCnt=0;
            lastCnt=rowLength;
        }else{
            fullCnt = ub_length / columnLength;
            lastCnt = rowLength - fullCnt * (loopCnt - 1);
        }
        //already aligned
        fullTileLength = fullCnt * columnLength;
        lastTileLength = lastCnt * columnLength;
    }else if(key==3){
        loopCnt = (totalLengthAligned + ub_length - 1) / ub_length;
        UB_SHARING_NUM=4;
        ub_length = ( (ub_size / typeSize / UB_SHARING_NUM) / ALIGN_NUM * ALIGN_NUM ) - ALIGN_NUM;
        fullTileLength = ub_length;
        lastTileLength = ( totalLengthAligned - fullTileLength * (loopCnt - 1) + ALIGN_NUM - 1 ) / ALIGN_NUM * ALIGN_NUM;
        if(loopCnt==1){
            fullTileLength=0;
        }
    }else{
        loopCnt = (totalLengthAligned + ub_length - 1) / ub_length;
        fullTileLength = ub_length;
        lastTileLength = ( totalLengthAligned - fullTileLength * (loopCnt - 1) + ALIGN_NUM - 1 ) / ALIGN_NUM * ALIGN_NUM;
        if(loopCnt==1){
            fullTileLength=0;
        }
    }

    
    tiling.set_totalLengthAligned(totalLengthAligned);
    tiling.set_matrixNum(matrixNum);
    tiling.set_matrixSize(matrixSize);
    tiling.set_rowLength(rowLength);
    tiling.set_columnLength(columnLength);
    tiling.set_diagVal(diagVal);

    tiling.set_loopCnt(loopCnt);
    tiling.set_fullTileLength(fullTileLength);
    tiling.set_lastTileLength(lastTileLength);
    tiling.set_fullCnt(fullCnt);
    tiling.set_lastCnt(lastCnt);

    tiling.set_alignNum(ALIGN_NUM);
    tiling.set_typeSize(typeSize);
    
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    context->SetTilingKey(key);
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
/*	
    std::cout<<"Key is "<<key<<std::endl;
    std::cout<<"Diagonal is "<<diagVal<<std::endl;
    std::cout<<"MatrixNum is "<<matrixNum<<std::endl;
    std::cout<<"MatrixSize is "<<matrixSize<<std::endl;
    std::cout<<"rowLength is "<<rowLength<<std::endl;
    std::cout<<"columnLength is "<<columnLength<<std::endl;
    std::cout<<"FullTileLength is "<<fullTileLength<<std::endl;
    std::cout<<"LastTileLength is "<<lastTileLength<<std::endl;
*/
    return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class Triu : public OpDef {
public:
    explicit Triu(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("diagonal").AttrType(OPTIONAL).Int(0);

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");

    }
};

OP_ADD(Triu);
}
